import torch
from tqdm import tqdm
from torch_geometric.data import Data, InMemoryDataset
import pandas as pd
from sklearn.preprocessing import LabelEncoder
import time
import torch_geometric.utils
from torch_geometric.data import HeteroData
from torch_geometric.loader import LinkNeighborLoader
import torch
import networkx as nx
import matplotlib.pyplot as plt

import torch_geometric.transforms as T
from torch_geometric.nn import HGTConv, Linear
from origin_data import MyData

from load_data import alarms_df


class YooChooseBinaryDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super(YooChooseBinaryDataset, self).__init__(root, transform, pre_transform)
    @property
    def raw_file_names(self):
        return []

    def download(self):
        pass

    @property
    def processed_file_names(self):
        return ["yoochoose_click_binary_1M_sess_dataset"]

    def process(self):
        data_list = []
        data = MyData(alarms_df).data
        data = T.ToUndirected()(data)
        train_data, val_data, test_data = transform(data.to_homogeneous())
        print("Test_data",test_data)
        data_list.append(data)
        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])

transform = T.RandomLinkSplit(
    num_val=0.1,
    num_test=0.1,
    disjoint_train_ratio=0.0,
    neg_sampling_ratio=2.0,
    add_negative_train_samples=False,
    is_undirected=True,
    # split_labels=True,
    edge_types=[
        ("alarm", "on", "host"),
        ("alarm", "to", "bussiness_tree"),
        ("host", "belongsto", "bussiness_tree")
    ]
)


my_dataset = YooChooseBinaryDataset(root="data",transform=transform)



# print(my_dataset[0].x)







